#pragma once
#include "CSCMatrix.hpp"
#include "CSRMatrix.hpp"
#include <assert.h>
#include <unordered_map>

extern "C" {
#include <immintrin.h>
}

bool csr_equal(CSRMatrix m1, CSRMatrix m2) {
    assert(m1.rowCount_ == m2.rowCount_ && m1.colCount_ == m2.colCount_);
    for (int i = 0; i < m1.rowCount_; i++) {
        for (int j = 0; j < m2.colCount_; j++) {
            if (fabs(m1.get(i, j) - m2.get(i, j)) > 1e-5) {
                std::cout << "\n"
                          << m1.get(i, j) << " != " << m2.get(i, j) << std::endl;
                return false;
            }
        }
    }
    return true;
}

void sp_mm_v1(CSRMatrix &a, CSCMatrix &b, CSRMatrix &c) {
    if (a.colCount_ != b.rowCount_) {
        std::cout << "Cannot multiply: Left matrix column count and right matrix row count don't "
                     "match.\n";
    }

    double tmp = 0;
    for (int i = 0; i < a.rowCount_; i++) {
        for (int j = 0; j < b.colCount_; j++) {
            tmp = 0;

            for (int k = 0; k < a.colCount_; k++) {
                tmp = tmp + a.get(i, k) * b.get(k, j);
            }

            c.set(tmp, i, j);
        }
    }
}

void sp_mm_v2(CSRMatrix &a, CSCMatrix &b, CSRMatrix &c) {
    if (a.colCount_ != b.rowCount_) {
        std::cout << "Cannot multiply: Left matrix column count and right matrix row count don't "
                     "match.\n";
    }

    int num_rows_A = a.rowCount_;
    int num_cols_B = b.colCount_;

    c.rowOffset_.resize(num_rows_A + 1);
    c.data_.clear();
    c.colInd_.clear();

    c.rowOffset_[0] = 0;
    for (int i = 0; i < num_rows_A; i++) {
        for (int j = 0; j < num_cols_B; j++) {
            double sum = 0.0;
            int index_A = a.rowOffset_[i];
            int end_A = a.rowOffset_[i + 1];
            int index_B = b.colOffset_[j];
            int end_B = b.colOffset_[j + 1];

            while (index_A < end_A && index_B < end_B) {
                if (a.colInd_[index_A] == b.rowInd_[index_B]) {
                    sum += a.data_[index_A] * b.data_[index_B];
                    index_A++;
                    index_B++;
                } else if (a.colInd_[index_A] < b.rowInd_[index_B]) {
                    index_A++;
                } else {
                    index_B++;
                }
            }

            if (sum != 0.0) {
                c.data_.push_back(sum);
                c.colInd_.push_back(j);
            }
        }
        c.rowOffset_[i + 1] = c.data_.size();
    }
}

void sp_mm_v3(CSRMatrix &a, CSCMatrix &b, CSRMatrix &c) {
    if (a.colCount_ != b.rowCount_) {
        std::cout << "Cannot multiply: Left matrix column count and right matrix row count don't "
                     "match.\n";
    }

    int num_rows_A = a.rowCount_;
    int num_cols_B = b.colCount_;

    c.rowOffset_.resize(num_rows_A + 1);
    c.data_.clear();
    c.colInd_.clear();

    c.rowOffset_[0] = 0;
    for (int i = 0; i < num_rows_A; i++) {
        for (int j = 0; j < num_cols_B; j++) {
            double sum = 0.0;
            int index_A = a.rowOffset_[i];
            int end_A = a.rowOffset_[i + 1];
            if (index_A - end_A == 0)
                break;
            int index_B = b.colOffset_[j];
            int end_B = b.colOffset_[j + 1];

            while (index_A < end_A && index_B < end_B) {
                if (a.colInd_[index_A] == b.rowInd_[index_B]) {
                    sum += a.data_[index_A] * b.data_[index_B];
                    index_A++;
                    index_B++;
                } else if (a.colInd_[index_A] < b.rowInd_[index_B]) {
                    index_A++;
                } else {
                    index_B++;
                }
            }

            if (sum != 0.0) {
                c.data_.push_back(sum);
                c.colInd_.push_back(j);
            }
        }
        c.rowOffset_[i + 1] = c.data_.size();
    }
}
